#include <mytbf.h>
#include <unistd.h>
#include <fcntl.h>
#include <pthread.h>
#include <sys/types.h>
#include <time.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>

// 令牌桶
struct mytbf_st {
    int cps;    // 速率
    int burst;  // 令牌个数上限
    int token;  // 令牌个数
    int pos;    // 在令牌桶数组里的位置
    pthread_mutex_t mut;    // 保护的此令牌桶的数据
    pthread_cond_t cond;    // 保护的此令牌桶的数据
};

static struct mytbf_st* job[MYTBF_MAX] = {NULL}; // 令牌桶数组，存放令牌桶
static pthread_mutex_t mut_job = PTHREAD_MUTEX_INITIALIZER; // 保护令牌桶数组
static pthread_cond_t cond_job = PTHREAD_COND_INITIALIZER;  // 保护令牌桶数组
static pthread_once_t init_once = PTHREAD_ONCE_INIT;    // 保证当前模块加载\初始化一次
static pthread_t tid;   // 线程thr_alrm的线程编号

// 对每个令牌桶生成令牌
static void* thr_alrm(void* arg) {
    while (1) {
        pthread_mutex_lock(&mut_job);

        int i = 0;
        for (i = 0; i < MYTBF_MAX; i++) {
            if (job[i] != NULL) {
                pthread_mutex_lock(&job[i]->mut);
                job[i]->token += job[i]->cps;
                if (job[i]->token > job[i]->burst) {
                    job[i]->token = job[i]->burst;
                }
                pthread_cond_broadcast(&job[i]->cond);  // 通知正在等待令牌桶job[i]的线程这个令牌桶的token增加了
                pthread_mutex_unlock(&job[i]->mut);
            }
        }

        pthread_mutex_unlock(&mut_job);
        sleep(1);
    }
}

static void module_unlocd(void) {
    pthread_cancel(tid);
    pthread_join(tid, NULL);
    int i = 0;
    for (i = 0; i < MYTBF_MAX; i++) {
        if (job[i] != NULL) {
            free(job[i]);
            job[i] = NULL;
        }
    }
}

static void module_load(void) {
    int ret = -1;
    ret = pthread_create(&tid, NULL, thr_alrm, NULL);
    if (ret == -1) {
        fprintf(stderr, "pthread_create error,%s\n", strerror(errno));
        exit(1);
    }

    atexit(module_unlocd);
}

// 取令牌桶编号(非线程安全)
static int get_free_pos_unlocked(void) {
    int i;
    for (i = 0; i < MYTBF_MAX; i++) {
        if (job[i] == NULL) {
            return i;
        }
    }
    return -1;
}

mytbf_t* mytbf_init(int cps, int burst) {
    pthread_once(&init_once, module_load);  // 指定module_load函数只执行一次
    module_load();

    struct mytbf_st* me = (struct mytbf_st*)malloc(sizeof(*me));
    if (me == NULL) {
        return NULL;
    }
    me->cps = cps;
    me->burst = burst;
    me->token = 0;
    pthread_mutex_lock(&mut_job);
    me->pos = get_free_pos_unlocked();
    if (me->pos < 0) {
        pthread_mutex_unlock(&mut_job);
        free(me);
        return NULL;
    }
    pthread_mutex_init(&me->mut, NULL);
    pthread_cond_init(&me->cond, NULL);
    job[me->pos] = me; 
    pthread_mutex_unlock(&mut_job);
    return me;
}

int mytbf_fetch_token(mytbf_t* tbf, int size) {
    struct mytbf_st* me = tbf;
    pthread_mutex_lock(&me->mut);
    while (me->token <= 0) {
        pthread_cond_wait(&me->cond, &me->mut); // 等待当前令牌桶tbf产生令牌
    }
    int n = ((me->token < size) ? (me->token) : (size));
    me->token -= n;
    pthread_mutex_unlock(&me->mut);
    return n;
}

int mytbf_return_token(mytbf_t* tbf, int size) {
    struct mytbf_st* me = tbf;
    pthread_mutex_lock(&me->mut);
    me->token += size;
    if (me->token > me->burst) {
        me->token = me->burst;
    }
    pthread_cond_broadcast(&me->cond);  // 通知在等待tbf的线程，tbf的token增加了
    pthread_mutex_unlock(&me->mut);
    return 0;
} 

int mytbf_destroy(mytbf_t* tbf) {
    struct mytbf_st* me = tbf;
    pthread_mutex_lock(&mut_job);
    job[me->pos] = NULL;
    pthread_mutex_destroy(&me->mut);
    pthread_cond_destroy(&me->cond);
    pthread_mutex_unlock(&mut_job);
    return 0;
}